#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This is a modified version of msoffice-product-key-decoder by
# prometheus235@gmail.com
# http://code.google.com/p/msoffice-product-key-decoder/
# Apache License 2.0 <http://www.apache.org/licenses/LICENSE-2.0>
#
# Modified by Nicholas Wilde 10/11/2012

import sys, string, math, collections
import _winreg as wr
from core import regutils

code_len = 25           # encoded key length (user-readable key)
bin_len = 15            # binary key length
regkey_idx = 52         # start of key in DPID for 2003, 2007
regkey_idx_2010 = 0x328 # start in DPID for 2010

generic_b24chrs = '0123456789ABCDEFGHIJKLMN'
b24chrs = 'BCDFGHJKMPQRTVWXY2346789'
reg_mso_root = r'Software\Microsoft\Office'
reg_msw_root = r'Software\Microsoft\Windows NT\CurrentVersion'
        
def b24encode(input, outlen=None, chrmap=None):

  # The number of coded characters to actually generate can't be 
  # determined soley from the input length, but we can guess.
  if outlen is None:
    # each base24 code char takes ~4.585 bits (ln24 / ln2)
    outlen = int(math.ceil(8*len(input) / 4.585))
      
  # Use default character mapping [0-9A-N] if none provided
  if chrmap is None:
    chrmap = generic_b24chrs
      
  input = [ord(i) for i in input[::-1]]
  '''
  # takes less memory (does it piecewise), but more complex
  decoded = []
  for i in range(0,encoded_chars + 1)[::-1]:
    r = 0
    for j in range(0,15)[::-1]:
      r = (r * 256) ^ input[j]
      input[j] = r / 24
      r = r % 24
    
    print b24chrs[r]
    decoded = decoded.append(b24chrs[r])
  
  return decoded[::-1]
  '''
  
  # simple, but eats a ton of memory and probably time if the 
  # encoded string is large
  enc = 0
  for i in input:
    enc = enc * 256 + i
      
  dec = []
  for i in range(outlen):
    dec.append(chrmap[enc % 24])
    enc = enc // 24
      
  dec.reverse()
  return ''.join(dec)

def b24decode(input, chrmap=None):

  # Use default character mapping [0-9A-N] if none provided
  if chrmap is None:
    chrmap = generic_b24chrs
  
  # clean invalid characters from input (e.g. '-' (dashes) in product key)
  # and map to \x00 through \x23.
  rmchrs = []
  for i in xrange(256):
    if not chr(i) in chrmap:
      rmchrs.append(chr(i))
  tt = string.maketrans(chrmap, ''.join([chr(i) for i in xrange(24)]))
  input = input.translate(tt, ''.join(rmchrs))
      
  encnum = 0
  for cc in input:
    encnum *= 24
    encnum += ord(cc)
      
  enc = []
  while encnum:
    enc.append(encnum % 256)
    encnum = encnum // 256
  
  return ''.join([chr(i) for i in enc])
    
def msKeyDecode(regkey, version=None):
  '''
  Decodes a registry key value, by extracting product key 
  from bytes 52-66 and decoding.
  
  Office 2010 (14.0) appears to store the key at 0x328 to 0x337 in 
  DPID.  The "Registration" tree is different (cluttered) versus 
  other versions, and the DPID value is (exactly) 7 times longer than 
  before (1148 bytes, up from 164).
  
  Tested with a 2010 full suite and a trial of Visio.
  
  Parameters:
  - regkey is a string containing the contents of "DigitalProductID"
  - version is the decimal version number given by the key directly 
  under the "Office" key root
  '''
  if version is None:
    version = 11 # (default 2003, 2007 appears to be compatible.)
  if float(version) < 14:
    enckey = regkey[regkey_idx:regkey_idx+bin_len]
  else:
    enckey = regkey[regkey_idx_2010:regkey_idx_2010+bin_len]

  deckey = b24encode(enckey, code_len, chrmap=b24chrs)
  return '-'.join(list(regutils.chunks(deckey,5)))

def getOfficeKeys():
  '''
  Scans local Microsoft Office registry keys for DigitalProductID values
  and encodes the binary data in base24.
  
  Note: The given "Name:" of Office 2010 products is incorrect
  (may just provide a single program name), though the Product Key 
  should be valid.
  '''
  
  prod_keys = []
  flags = regutils.getOSFlagList()
  for flag in flags:
    mso_root = wr.OpenKey(wr.HKEY_LOCAL_MACHINE, reg_mso_root, 0,
      wr.KEY_READ | flag)
    
    for subkey in regutils.SubKeys(mso_root):
      # subkey always observed to be a version number (11.0, 12.0, 14.0...)
      # where a DPID will be found.
      for sub2key in regutils.SubKeys(wr.OpenKey(mso_root, subkey, 0,
        wr.KEY_READ | flag)):
        # sub2key always observed to be "Registration" when DPID is found
        for sub3key in regutils.SubKeys(wr.OpenKey(wr.OpenKey(mso_root, subkey, 0,
          wr.KEY_READ | flag), sub2key, 0, wr.KEY_READ | flag)):
          # sub3key often some UUID ending in 0F:F1:CE
          dpid_found = False
          for keyvalue in regutils.KeyValues(wr.OpenKey(wr.OpenKey(wr.OpenKey(
            mso_root, subkey, 0, wr.KEY_READ | flag), sub2key, 0,
            wr.KEY_READ | flag), sub3key, 0, wr.KEY_READ | flag)):
            if keyvalue[0].lower() == 'digitalproductid':
              dpid_found = True
              dpid = keyvalue
            if keyvalue[0].lower() == 'productname':
              name = keyvalue
            if keyvalue[0].lower() == 'productid':
              pid = keyvalue
          if dpid_found:
            prod_keys.append(collections.OrderedDict([
              ('Name:', name[1]),
              ('Product ID:', pid[1]),
              ('CD Key:', msKeyDecode(dpid[1], subkey))
            ]))
  if len(prod_keys) is 0:
    prod_keys = None
  return prod_keys

def getWinKey():
  """Get Windows Key"""
  win_key = []
  dpid_found = False
  flag = regutils.getOSFlag()
  msw_root = wr.OpenKey(wr.HKEY_LOCAL_MACHINE, reg_msw_root, 0,
    wr.KEY_READ | flag)  
  for keyvalue in regutils.KeyValues(msw_root):
    if keyvalue[0].lower() == 'digitalproductid':
      dpid_found = True
      dpid = keyvalue
    if keyvalue[0].lower() == 'productname':
      name = keyvalue
    if keyvalue[0].lower() == 'productid':
      pid = keyvalue
    if keyvalue[0].lower() == 'csdversion':
      csdv = keyvalue
    if keyvalue[0].lower() == 'registeredowner':
      ro = keyvalue
    if keyvalue[0].lower() == 'registeredorganization':
        rorg = keyvalue
  if dpid_found:
    win_key.append(collections.OrderedDict([
      ('Name:', name[1]),
      ('Product ID:', pid[1]),
      ('CD Key:', msKeyDecode(dpid[1])),
      ('Computer Name:', regutils.getComputerName()),
      ('Registered Owner:', ro[1]),
      ('Registered Organization:', rorg[1])
    ]))
  if len(win_key) is 0:
    win_key = None
  return win_key